#include "stdafx.h"
#include "SocketServerTestSuite.h"

namespace CO { namespace SpiFoundation { namespace RoomFW {

class TestSocketServer : public CO::SpiFoundation::RoomFW::SocketServer
{
public:
	TestSocketServer(void);
	bool waitEmpty(DWORD ms);

protected:
	virtual bool handleConnect(CO::SpiFoundation::RoomFW::SimpleSocket::SocketPtr client);
	virtual bool handleDisconnect(CO::SpiFoundation::RoomFW::SimpleSocket::SocketPtr client);
	virtual bool handleMessage(CO::SpiFoundation::RoomFW::SimpleSocket::SocketPtr client, const std::string& message);

private:

	CO::SpiFoundation::RoomFW::SimpleEvent empty_;
	std::set<CO::SpiFoundation::RoomFW::SimpleSocket::SocketPtr> socks_;
};

TestSocketServer::TestSocketServer(void)
{
	empty_.create(true, true);
}

bool TestSocketServer::waitEmpty(DWORD ms)
{
	return empty_.wait(ms);
}

bool TestSocketServer::handleConnect(CO::SpiFoundation::RoomFW::SimpleSocket::SocketPtr client)
{
	socks_.insert(client);
	empty_.reset();
	client->sendMessage("connected");
	return true;
}

bool TestSocketServer::handleDisconnect(CO::SpiFoundation::RoomFW::SimpleSocket::SocketPtr client)
{
	socks_.erase(client);
	if (socks_.empty())
	{
		empty_.set();
	}
	return true;
}

bool TestSocketServer::handleMessage(CO::SpiFoundation::RoomFW::SimpleSocket::SocketPtr client, const std::string& message)
{
	client->sendMessage(message);
	return true;
}

SocketServerTestSuite::SocketServerTestSuite(void)
{
	TEST_ADD(SocketServerTestSuite::all)
}

SocketServerTestSuite::~SocketServerTestSuite(void)
{
}

void SocketServerTestSuite::setup()
{
	// Initialize Winsock
	WSADATA wsaData;
	int iResult = WSAStartup(MAKEWORD(2,2), &wsaData);
	if (0 != iResult)
	{
		std::cout << "Error " << iResult << " at WSAStartup()" << std::endl;
	}
}

void SocketServerTestSuite::tear_down()
{
	// Shutdown Winsock
	int iResult = WSACleanup();
	if (SOCKET_ERROR == iResult)
	{
		std::cout << "Error " << WSAGetLastError() << " at WSAStartup()" << std::endl;
	}
}

std::string testMsg(void)
{
	std::string retVal;
	retVal.reserve(1024);
	for (int i = 0; i < 1024; ++i)
	{
		retVal.push_back(static_cast<char>(CO::SpiFoundation::RoomFW::WindowsUtility::Instance().rand(0, CHAR_MAX)));
	}
	return retVal;
}

void SocketServerTestSuite::all(void)
{
	// Start a server.
	TestSocketServer server;
	TEST_ASSERT(server.start(2001, "127.0.0.1"));

	// Run the following twice to ensure server is cleaning up.  It shouldn't be
	// able to handle more than 62 clients at a time.
	for (int i = 0; i < 2; ++i)
	{
		// Connect 50 clients.
		std::vector<CO::SpiFoundation::RoomFW::SimpleSocket::SocketPtr> clients;
		for (int j = 0; j < 50; ++j)
		{
			CO::SpiFoundation::RoomFW::SimpleSocket::SocketPtr client(new CO::SpiFoundation::RoomFW::SimpleSocket());
			TEST_ASSERT(client->connect(2001, "127.0.0.1"));
			std::string conMsg;
			TEST_ASSERT(client->receiveMessage(conMsg));
			TEST_ASSERT(conMsg == "connected");
			std::string sndMsg = testMsg();
			TEST_ASSERT(client->sendMessage(sndMsg));
			std::string rcvMsg;
			TEST_ASSERT(client->receiveMessage(rcvMsg));
			TEST_ASSERT(rcvMsg == sndMsg);
			clients.push_back(client);
		}

		// Randomly re-test and delete.
		while (!clients.empty())
		{
			const int index = CO::SpiFoundation::RoomFW::WindowsUtility::Instance().rand(0, clients.size() - 1);
			CO::SpiFoundation::RoomFW::SimpleSocket::SocketPtr client = clients[index];
			std::string sndMsg = testMsg();
			TEST_ASSERT(client->sendMessage(sndMsg));
			std::string rcvMsg;
			TEST_ASSERT(client->receiveMessage(rcvMsg));
			TEST_ASSERT(rcvMsg == sndMsg);
			clients.erase(clients.begin() + index);
		}

		// Make sure everything disconnected.
		TEST_ASSERT(server.waitEmpty(2000));
	}
}

}}}
